from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 诊断相关基础数据 ==========
DIAG_TYPES = ['西医诊断', '中医诊断', '入院诊断', '出院诊断', '术前诊断', '术后诊断']
DIAG_STATUSES = ['确诊', '疑似', '待查', '排除']
CONFIRMED_DIAG_MARKS = ['是', '否']
MAINDIAG_MARKS = ['是', '否']
IS_COMPLICATIONS = ['是', '否']
VISIT_TYPES = ['门诊', '急诊', '住院']
DEPARTMENTS = ['肿瘤科', '放疗科', '化疗中心', '内科', '外科', '中医科']

# 常见癌症诊断（ICD-10编码）
CANCER_DIAGNOSES = [
    ("C34.901", "肺恶性肿瘤", "西医诊断"),
    ("C50.902", "乳腺恶性肿瘤", "西医诊断"),
    ("C16.901", "胃恶性肿瘤", "西医诊断"),
    ("C18.901", "结肠恶性肿瘤", "西医诊断"),
    ("C22.001", "肝细胞癌", "西医诊断"),
    ("C25.901", "胰腺恶性肿瘤", "西医诊断"),
    ("C73.x00", "甲状腺恶性肿瘤", "西医诊断"),
    ("C91.101", "急性淋巴细胞白血病", "西医诊断"),
    ("C92.901", "慢性粒细胞白血病", "西医诊断"),
    ("C85.903", "非霍奇金淋巴瘤", "西医诊断")
]

# 中医诊断数据
TCM_DIAGNOSES = [
    ("BNX010", "肺癌病", "BNX01001", "气阴两虚证"),
    ("BNX020", "乳岩病", "BNX02001", "肝郁气滞证"),
    ("BNX030", "胃癌病", "BNX03001", "脾胃虚弱证"),
    ("BNX040", "肝癌病", "BNX04001", "肝郁脾虚证")
]

HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 15 if date_obj.day % 2 == 0 else 8


# ========== 核心生成函数 ==========
def generate_diagnosis_records():
    sql_statements = []
    record_count = 0

    patient_id_counter = 2001
    visit_sn_counter = 2001
    diag_id_counter = 50001

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1
            patient_id = f"PT{patient_id_counter}"
            patient_id_counter += 1

            visit_sn = f"VIS{visit_sn_counter}"
            visit_sn_counter += 1

            # 随机选择医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 基础患者信息
            medical_record_no = f"MR{random.randint(100000, 999999)}"
            inpatient_no = f"IP{random.randint(100000, 999999)}" if random.random() > 0.3 else None
            outpatient_no = f"OP{random.randint(100000, 999999)}" if random.random() > 0.3 else None
            visit_card_no = f"VC{random.randint(100000, 999999)}" if random.random() > 0.5 else None
            name = fake.name()

            # 诊断信息
            diag_id = f"DIAG{diag_id_counter}"
            diag_id_counter += 1

            diag_serial_number = str(random.randint(1, 10))
            maindiag_mark = random.choice(MAINDIAG_MARKS)
            confirmed_diag_mark = random.choice(CONFIRMED_DIAG_MARKS)
            is_complication = random.choice(IS_COMPLICATIONS)

            # 随机选择西医或中医诊断
            if random.random() > 0.7:  # 30%概率生成中医诊断
                diag_type = "中医诊断"
                tcm_disease_code, tcm_disease_name, tcm_symptom_code, tcm_symptom_name = random.choice(TCM_DIAGNOSES)
                diag_name = f"{tcm_disease_name}-{tcm_symptom_name}"
                diag_code = tcm_disease_code
            else:
                diag_type = random.choice(["西医诊断", "入院诊断", "出院诊断"])
                diag_code, diag_name, _ = random.choice(CANCER_DIAGNOSES)
                tcm_disease_code = tcm_disease_name = tcm_symptom_code = tcm_symptom_name = None

            diag_source = random.choice(["HIS系统", "电子病历", "门诊系统"])
            diag_status = random.choice(DIAG_STATUSES)
            visit_type = random.choice(VISIT_TYPES)

            # 时间相关字段
            diag_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d"),
                end_date=(datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1))
            ).strftime('%Y-%m-%d %H:%M:%S')

            record_datetime = diag_datetime
            yy_collection_datetime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            yy_etl_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # 科室和医生信息
            create_diagnose_dept_name = random.choice(DEPARTMENTS)
            create_diagnose_dept_code = f"DEPT{random.randint(10, 99)}"
            diag_doctor_name = fake.name()
            diag_doctor_no = f"DOC{random.randint(100, 999)}"

            # 其他字段
            hospitalization_times = str(random.randint(1, 5)) if visit_type == '住院' else None
            visit_times = str(random.randint(1, 10))
            diag_explanation = f"患者因{random.choice(['体检发现', '症状就诊'])}，经{random.choice(['影像学检查', '病理检查', '实验室检查'])}确诊为{diag_name}"

            # 生成yy_record_md5 (模拟MD5)
            yy_record_md5 = f"md5_{random.getrandbits(128):032x}"
            yy_batch_time = day
            yy_record_batch_id = f"BATCH{day}_{i}"

            data = {
                'confirmed_diag_mark': confirmed_diag_mark,
                'create_diagnose_dept_code': create_diagnose_dept_code,
                'create_diagnose_dept_name': create_diagnose_dept_name,
                'diag_code': diag_code,
                'diag_datetime': diag_datetime,
                'diag_doctor_name': diag_doctor_name,
                'diag_doctor_no': diag_doctor_no,
                'diag_explanation': diag_explanation,
                'diag_id': diag_id,
                'diag_name': diag_name,
                'diag_serial_number': diag_serial_number,
                'diag_source': diag_source,
                'diag_status': diag_status,
                'diag_type': diag_type,
                'extend_data1': None,
                'extend_data2': None,
                'from_table': 'EMR_DIAGNOSIS',
                'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                'hospital_code': hospital_code,
                'hospital_name': hospital_name,
                'hospitalization_times': hospitalization_times,
                'inpatient_no': inpatient_no,
                'is_complication': is_complication,
                'maindiag_mark': maindiag_mark,
                'medical_record_no': medical_record_no,
                'name': name,
                'outpatient_no': outpatient_no,
                'patient_id': patient_id,
                'patient_id_old': f"OLD_{patient_id}",
                'record_datetime': record_datetime,
                'record_status': 1,
                'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'tcm_disease_code': tcm_disease_code,
                'tcm_disease_name': tcm_disease_name,
                'tcm_symptom_code': tcm_symptom_code,
                'tcm_symptom_name': tcm_symptom_name,
                'visit_card_no': visit_card_no,
                'visit_sn': visit_sn,
                'visit_times': visit_times,
                'visit_type': visit_type,
                'yy_collection_datetime': yy_collection_datetime,
                'yy_record_md5': yy_record_md5,
                'yy_upload_status': 0,
                'yy_etl_time': yy_etl_time,
                'yy_upload_time': None,
                'yy_batch_time': yy_batch_time,
                'yy_record_batch_id': yy_record_batch_id,
                'yy_backfill_time': None,
                'yy_backfill_status': None,
                'branch_code': branch_code,
                'branch_name': branch_name,
                'date_for_partition': diag_datetime
            }

            sql = _generate_sql('b02_3', data)
            sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_diagnosis_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b02_3')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b02_3_diagnosis_records.sql")